#!/usr/bin/env python3

import argparse
import json
import os
import re
import subprocess
import sys
import tempfile
from collections import OrderedDict


builddir = "@builddir@"


# Parse command-line arguments.
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output",
        help="JSON file to write benchmark results to")
    parser.add_argument("--verify", action='store_true',
        help="verify parse results on full data files")
    parser.add_argument("--quickverify", action='store_true',
        help="verify parse results on small data samples")

    args = parser.parse_args()
    if args.output == None and not (args.verify or args.quickverify):
        sys.exit("no output file!")

    return args


def merge(combined, input_gcc, input_clang, output):
    for i in [input_gcc, input_clang]:
        with open(i) as f:
            contents = json.load(f)
        for bench in contents['benchmarks']:
            name = bench['name']
            combined.setdefault(name, {})
            combined[name]['name'] = name
            if 'cpu_time' in bench:
                combined[name]['cpu_time'] = bench['cpu_time']
            if 'error_occurred' in bench:
                combined[name]['error'] = bench['error_occurred']

    filtered = []
    for k, v in combined.items():
        if not 'error' in v:
            filtered.append(v)

    with open(output, 'w') as f:
        json.dump({'benchmarks': filtered}, f, indent = 2)


def get_sym_size(fname, pattern):
    n = 0
    p = subprocess.run(["nm", "-S", "-C", fname], stdout=subprocess.PIPE, check=True)
    for m in re.finditer(pattern, p.stdout):
        n += int(m.group(1), 16)
    return round(n / 1024.0, 1)


def bench_size():
    benchmarks = OrderedDict()

    for f in sorted(os.listdir(os.path.join(builddir, "gen", "re2c"))):
        if not f.endswith("-gcc.o"): continue

        for compiler in ["gcc", "clang"]:
            fname = re.search("([^/]+)-gcc.o$", f).group(1) + "-" + compiler + ".o"
            name = re.search("^.*__(.*)-gcc.o$", f).group(1).replace("_", "-") + "-" + compiler

            benchmarks[name + "_re2c-simple"] = { "bin_size": get_sym_size(
                os.path.join(builddir, "gen", "re2c", fname),
                b"([0-9a-f]{16}) [tT] .*::(lex|bench|fill)_simple") }

            benchmarks[name + "_re2c-buffered-eof"] = { "bin_size": get_sym_size(
                os.path.join(builddir, "gen", "re2c", fname),
                b"([0-9a-f]{16}) [tT] .*::(lex|bench|fill)_buffered_eof") }

            benchmarks[name + "_re2c-buffered-scc"] = { "bin_size": get_sym_size(
                os.path.join(builddir, "gen", "re2c", fname),
                b"([0-9a-f]{16}) [tT] .*::(lex|bench|fill)_buffered_scc") }

            benchmarks[name + "_ragel-simple"] = { "bin_size": get_sym_size(
                os.path.join(builddir, "gen", "ragel", fname),
                b"([0-9a-f]{16}) [tT] .*::(lex|bench_simple)") }

            benchmarks[name + "_ragel-buffered"] = { "bin_size": get_sym_size(
                os.path.join(builddir, "gen", "ragel", fname),
                b"([0-9a-f]{16}) [tT] .*::(lex|bench_buffered)") }

    return benchmarks


def main():
    args = parse_args()

    if args.verify:
        subprocess.run([os.path.join(builddir, "test")], check=True)
    elif args.quickverify:
        subprocess.run([os.path.join(builddir, "test"), "--quick"], check=True)
    else:
        tmp = tempfile.TemporaryDirectory()
        fgcc = os.path.join(tmp.name, "results_gcc.json")
        fclang = os.path.join(tmp.name, "results_clang.json")

        benchmarks = bench_size()

        subprocess.run([
                os.path.join(builddir, "bench-gcc"),
                "--benchmark_out_format=json",
                "--benchmark_out=" + fgcc
            ], check=True)

        subprocess.run([
                os.path.join(builddir, "bench-clang"),
                "--benchmark_out_format=json",
                "--benchmark_out=" + fclang
            ], check=True)

        merge(benchmarks, fgcc, fclang, args.output)

        os.remove(fgcc)
        os.remove(fclang)


if __name__ == "__main__":
    main()
